{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [FlashRank](https://github.com/PrithivirajDamodaran/FlashRank) 是一个超轻量、超快速的库，用于向现有的搜索和检索流程中添加重新排名功能。它基于最先进的交叉编码器技术，这使得它能够在处理大量数据时保持高效率。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "%pip install --upgrade --quiet flashrank  \n",
    "%pip install --upgrade --quiet faiss  \n",
    "如果只需要cpu版本  \n",
    "%pip install --upgrade --quiet faiss_cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "##打印文本内容\n",
    "def pretty_print_docs(docs):\n",
    "    print(\n",
    "        f\"\\n{'-' * 100}\\n\".join(\n",
    "            [\n",
    "                f\"Document {i+1}:\\n\\n{d.page_content}\\nMetadata: {d.metadata}\"\n",
    "                for i, d in enumerate(docs)\n",
    "            ]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 我们开始初始化一个简单的向量存储检索器\n",
    "\n",
    "以[baichuan2](https://arxiv.org/pdf/2309.10305v2.pdf)论文为例 ， 我们可以设置检索器以检索大量（20个）文档。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "documents = PyPDFLoader(\"https://arxiv.org/pdf/2309.10305.pdf\").load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
    "texts = text_splitter.split_documents(documents[:15])\n",
    "for idx, text in enumerate(texts):\n",
    "    text.metadata[\"id\"] = idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "compared to Baichuan 1-7B. Specifically, Baichuan\n",
      "2 is optimized to improve performance on math\n",
      "and code problems. On the GSM8K (Cobbe\n",
      "et al., 2021) and HumanEval (Chen et al., 2021)\n",
      "evaluations, Baichuan 2 nearly doubles the results\n",
      "of the Baichuan 1. In addition, Baichuan 2 also\n",
      "demonstrates strong performance on medical and\n",
      "legal domain tasks. On benchmarks such as\n",
      "MedQA (Jin et al., 2021) and JEC-QA (Zhong\n",
      "et al., 2020), Baichuan 2 outperforms other open-\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 14}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "to Baichuan 1-7B, Baichuan 2-7B-Base shows an\n",
      "improvement of nearly 10 points. In the medical\n",
      "field, Baichuan 2-7B-Base outperforms models\n",
      "like ChatGLM 2-6B and LLaMA 2-7B, showing\n",
      "significant improvement over Baichuan 1-7B as\n",
      "well.\n",
      "Similarly, Baichuan 2-13B-Base surpasses\n",
      "models other than GPT-4 in the field of Chinese\n",
      "law. In the medical domain, Baichuan 2-13B-\n",
      "Base outperforms models such as XVERSE-13B\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 8, 'id': 89}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "parameters, trained from scratch, on 2.6 trillion\n",
      "tokens. Baichuan 2 matches or outperforms\n",
      "other open-source models of similar size on\n",
      "public benchmarks like MMLU, CMMLU,\n",
      "GSM8K, and HumanEval. Furthermore,\n",
      "Baichuan 2 excels in vertical domains such\n",
      "as medicine and law. We will release all\n",
      "pre-training model checkpoints to benefit the\n",
      "research community in better understanding\n",
      "the training dynamics of Baichuan 2.\n",
      "1 Introduction\n",
      "The field of large language models has witnessed\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 0, 'id': 3}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 4:\n",
      "\n",
      "to our knowledge is the largest to date, more than\n",
      "double that of Baichuan 1 (Baichuan, 2023b,a).\n",
      "With such a massive amount of training data,\n",
      "Baichuan 2 achieves significant improvements over\n",
      "Baichuan 1. On general benchmarks like MMLU\n",
      "(Hendrycks et al., 2021a), CMMLU (Li et al.,\n",
      "2023), and C-Eval (Huang et al., 2023), Baichuan\n",
      "2-7B achieves nearly 30% higher performance\n",
      "compared to Baichuan 1-7B. Specifically, Baichuan\n",
      "2 is optimized to improve performance on math\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 13}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 5:\n",
      "\n",
      "to advance collective knowledge in developing\n",
      "LLMs. Baichuan 2’s foundation models and\n",
      "chat models are available for both research and\n",
      "commercial use at https://github.com/\n",
      "baichuan-inc/Baichuan2\n",
      "2 Pre-training\n",
      "This section introduces the training procedure\n",
      "for the Baichuan 2 foundation models. Before\n",
      "diving into the model details, we first show the\n",
      "overall performance of the Baichuan 2 base models\n",
      "compared to other open or closed-sourced models\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 20}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 6:\n",
      "\n",
      "development and application of LLMs in specific\n",
      "languages, such as Chinese.\n",
      "In this technical report, we introduce Baichuan\n",
      "2, a series of large-scale multilingual language\n",
      "models. Baichuan 2 has two separate models,\n",
      "Baichuan 2-7B with 7 billion parameters and\n",
      "Baichuan 2-13B with 13 billion parameters. Both\n",
      "models were trained on 2.6 trillion tokens, which\n",
      "to our knowledge is the largest to date, more than\n",
      "double that of Baichuan 1 (Baichuan, 2023b,a).\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 12}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 7:\n",
      "\n",
      "Baichuan 2: Open Large-scale Language Models\n",
      "Aiyuan Yang, Bin Xiao, Bingning Wang, Borong Zhang, Chao Yin, Chenxu Lv, Da Pan\n",
      "Dian Wang, Dong Yan, Fan Yang, Fei Deng, Feng Wang, Feng Liu, Guangwei Ai\n",
      "Guosheng Dong, Haizhou Zhao, Hang Xu, Haoze Sun, Hongda Zhang, Hui Liu, Jiaming Ji\n",
      "Jian Xie, Juntao Dai, Kun Fang, Lei Su, Liang Song, Lifeng Liu, Liyun Ru, Luyao Ma\n",
      "Mang Wang, Mickel Liu, MingAn Lin, Nuolan Nie, Peidong Guo, Ruiyang Sun\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 0, 'id': 0}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 8:\n",
      "\n",
      "MedQA (Jin et al., 2021) and JEC-QA (Zhong\n",
      "et al., 2020), Baichuan 2 outperforms other open-\n",
      "source models, making it a suitable foundation\n",
      "model for domain-specific optimization.\n",
      "Additionally, we also released two chat\n",
      "models, Baichuan 2-7B-Chat and Baichuan 2-\n",
      "13B-Chat, optimized to follow human instructions.\n",
      "These models excel at dialogue and context\n",
      "understanding. We will elaborate on our\n",
      "approaches to improve the safety of Baichuan 2.\n",
      "By open-sourcing these models, we hope to enable\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 15}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 9:\n",
      "\n",
      "Baichuan 2 training data.\n",
      "Data processing : For data processing, we focus\n",
      "on data frequency and quality. Data frequency\n",
      "relies on clustering and deduplication. We built\n",
      "a large-scale deduplication and clustering system\n",
      "supporting both LSH-like features and dense\n",
      "embedding features. This system can cluster\n",
      "and deduplicate trillion-scale data within hours.\n",
      "Based on the clustering, individual documents,\n",
      "paragraphs, and sentences are deduplicated and\n",
      "scored. Those scores are then used for data\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 2, 'id': 26}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 10:\n",
      "\n",
      "350 iterations for all our chat models, resulting in\n",
      "Baichuan 2-7B-Chat and Baichuan 2-13B-Chat.\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 6, 'id': 68}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 11:\n",
      "\n",
      "LLaMA 2-13B 13.24\n",
      "Table 8: Toxigen results of Baichuan 2 foundation\n",
      "models compared with LLaMA 2.\n",
      "Inspired by BeaverTails Ji et al. (2023)9, we\n",
      "constructed the Baichuan Harmless Evaluation\n",
      "Dataset (BHED), covering 7 major safety\n",
      "categories of bias/discrimination ,insults/profanity ,\n",
      "illegal/unethical content ,physical health ,mental\n",
      "health ,financial privacy , and sensitive topics to\n",
      "evaluate the safety of our chat models.\n",
      "8https://github.com/microsoft/SafeNLP/\n",
      "tree/main\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 11, 'id': 108}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 12:\n",
      "\n",
      "XVERSE-13B 53.70 55.21 58.44 44.69 42.54 38.06 18.20 15.85\n",
      "Baichuan 1-13B-Base 52.40 51.60 55.30 49.69 43.20 43.01 26.76 11.5913B\n",
      "Baichuan 2-13B-Base 58.10 59.17 61.97 54.33 48.17 48.78 52.77 17.07\n",
      "Table 1: Overall results of Baichuan 2 compared with other similarly sized LLMs on general benchmarks. * denotes\n",
      "results derived from official websites.\n",
      "Figure 1: The distribution of different categories of\n",
      "Baichuan 2 training data.\n",
      "Data processing : For data processing, we focus\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 2, 'id': 25}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 13:\n",
      "\n",
      "Baichuan 1-7B model.\n",
      "In mathematics, Baichuan 2-13B-Base surpasses\n",
      "all models of the same size, approaching the level\n",
      "of GPT-3.5 Turbo. In the code domain, Baichuan\n",
      "2-13B-Base outperforms models like LLaMA 2-\n",
      "13B and XVERSE-13B. Baichuan 2-13B-Base\n",
      "demonstrates significant improvement compared to\n",
      "Baichuan 1-13B-Base.\n",
      "5.4 Multilingual\n",
      "We use Flores-101 (NLLB Team, 2022; Goyal\n",
      "et al., 2021; Guzmán et al., 2019) to evaluate\n",
      "multilingual ability. Flores-101 covers 101\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 9, 'id': 93}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 14:\n",
      "\n",
      "overall performance of the Baichuan 2 base models\n",
      "compared to other open or closed-sourced models\n",
      "in Table 1. We then describe our pre-training data\n",
      "and data processing methods. Next, we elaborate\n",
      "on the Baichuan 2 architecture and scaling results.\n",
      "Finally, we describe the distributed training system.\n",
      "2.1 Pre-training Data\n",
      "Data sourcing : During data acquisition, our\n",
      "objective is to pursue comprehensive data\n",
      "scalability and representativeness. We gather data\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 1, 'id': 21}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 15:\n",
      "\n",
      "and LLaMA 2-13B. Compared to Baichuan 1-\n",
      "13B-Base, Baichuan 2-13B-Base also exhibits\n",
      "remarkable improvement.\n",
      "5.3 Math and Code\n",
      "This section introduces the performance in\n",
      "mathematics and coding.\n",
      "We use GSM8K (Cobbe et al., 2021) (4-shot)\n",
      "andMATH (Hendrycks et al., 2021b) (4-shot) to\n",
      "evaluate the mathematical ability. MATH contains\n",
      "12,500 mathematical questions that are harder to\n",
      "be solved. To evaluate the model’s code ability, we\n",
      "report the scores in HumanEval (Chen et al., 2021)\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 9, 'id': 90}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 16:\n",
      "\n",
      "model could be used to generate harmful or\n",
      "misleading content. Although we try our best\n",
      "efforts to balance safety and utility, some safety\n",
      "measures may appear as over-cautions, affecting\n",
      "the model’s usability for certain tasks. We\n",
      "encourage users to make responsible and ethical\n",
      "use of Baichuan 2 models. Meanwhile, we will\n",
      "continue to optimize these issues and release\n",
      "updated versions in the future.\n",
      "References\n",
      "Yuvanesh Anand, Zach Nussbaum, Brandon Duderstadt,\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 13, 'id': 126}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 17:\n",
      "\n",
      "Baichuan 1-13B-Base 30.24 20.90 15.92 0.98 9.65 2.64 12.00 13.1913B\n",
      "Baichuan 2-13B-Base 30.61 22.11 17.27 2.39 14.17 11.58 14.53 16.09\n",
      "Table 7: The result of Baichuan 2 compared with other models on multilingual field.\n",
      "version from the SafeNLP project8, distinguishing\n",
      "neutral and hate types for the 13 minority groups,\n",
      "forming a 6-shot dataset consistent with the\n",
      "original Toxigen prompt format. Our decoding\n",
      "parameters use temperature 0.1 and top-p 0.9\n",
      "nucleus sampling.\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 11, 'id': 106}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 18:\n",
      "\n",
      "7Limitations and Ethical Considerations\n",
      "Like other large language models, Baichuan 2 also\n",
      "faces ethical challenges. It’s prone to biases and\n",
      "toxicity, especially given that much of its training\n",
      "data originates from the internet. Despite our best\n",
      "efforts to mitigate these issues using benchmarks\n",
      "like Toxigen (Hartvigsen et al., 2022), the risks\n",
      "cannot be eliminated, and toxicity tends to increase\n",
      "with model size. Moreover, the knowledge of\n",
      "Baichuan 2 models is static and can be outdated or\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 13, 'id': 124}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 19:\n",
      "\n",
      "GPUs, achieving a computational efficiency that\n",
      "exceeds 180 TFLOPS.\n",
      "3 Alignment\n",
      "Baichuan 2 also introduces the alignment\n",
      "procedure resulting in two chat models: Baichuan\n",
      "2-7B-Chat and Baichuan 2-13B-Chat. The\n",
      "alignment process of the Baichuan 2 encompasses\n",
      "two main components: Supervised Fine-Tuning\n",
      "(SFT) and Reinforcement Learning from Human\n",
      "Feedback (RLHF).\n",
      "3.1 Supervised Fine-Tuning\n",
      "During the supervised fine-tuning phase, we use\n",
      "human labelers to annotate prompts gathered from\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 6, 'id': 59}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 20:\n",
      "\n",
      "a variety of natural language tasks based\n",
      "on just a few examples of natural language\n",
      "instructions, reducing the need for extensive\n",
      "feature engineering. However, most powerful\n",
      "LLMs are closed-source or limited in their\n",
      "capability for languages other than English. In\n",
      "this technical report, we present Baichuan 2,\n",
      "a series of large-scale multilingual language\n",
      "models containing 7 billion and 13 billion\n",
      "parameters, trained from scratch, on 2.6 trillion\n",
      "tokens. Baichuan 2 matches or outperforms\n",
      "Metadata: {'source': 'https://arxiv.org/pdf/2309.10305.pdf', 'page': 0, 'id': 2}\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "embedding = OpenAIEmbeddings(model=\"text-embedding-ada-002\")\n",
    "retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={\"k\": 20})\n",
    "query = \"What is Baichuan2？\"\n",
    "docs = retriever.invoke(query)\n",
    "pretty_print_docs(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用FlashRank进行重排\n",
    "\n",
    "现在让我们用一个**ContextualCompressionRetriever**来包装我们的基础检索器，并使用**FlashrankRerank**作为压缩器。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running pairwise ranking..\n",
      "[9, 1, 19]\n"
     ]
    }
   ],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import FlashrankRerank\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "\n",
    "compressor = FlashrankRerank()\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What is Baichuan2？\"\n",
    ")\n",
    "print([doc.metadata[\"id\"] for doc in compressed_docs])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "重排序之后，top3 已经和之前的top3发生了很大的变化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "350 iterations for all our chat models, resulting in\n",
      "Baichuan 2-7B-Chat and Baichuan 2-13B-Chat.\n",
      "Metadata: {'id': 9, 'relevance_score': [0.9729050397872925, 0.06210247427225113]}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "to Baichuan 1-7B, Baichuan 2-7B-Base shows an\n",
      "improvement of nearly 10 points. In the medical\n",
      "field, Baichuan 2-7B-Base outperforms models\n",
      "like ChatGLM 2-6B and LLaMA 2-7B, showing\n",
      "significant improvement over Baichuan 1-7B as\n",
      "well.\n",
      "Similarly, Baichuan 2-13B-Base surpasses\n",
      "models other than GPT-4 in the field of Chinese\n",
      "law. In the medical domain, Baichuan 2-13B-\n",
      "Base outperforms models such as XVERSE-13B\n",
      "Metadata: {'id': 1, 'relevance_score': [0.9596770405769348, 0.08595798164606094]}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "a variety of natural language tasks based\n",
      "on just a few examples of natural language\n",
      "instructions, reducing the need for extensive\n",
      "feature engineering. However, most powerful\n",
      "LLMs are closed-source or limited in their\n",
      "capability for languages other than English. In\n",
      "this technical report, we present Baichuan 2,\n",
      "a series of large-scale multilingual language\n",
      "models containing 7 billion and 13 billion\n",
      "parameters, trained from scratch, on 2.6 trillion\n",
      "tokens. Baichuan 2 matches or outperforms\n",
      "Metadata: {'id': 19, 'relevance_score': [0.6193831562995911, 0.4349020719528198]}\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "chain = RetrievalQA.from_chain_type(llm=llm, retriever=compression_retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running pairwise ranking..\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'What is Baichuan2？',\n",
       " 'result': 'Baichuan 2 is a series of large-scale multilingual language models containing 7 billion and 13 billion parameters. These models were trained from scratch on 2.6 trillion tokens and are designed to perform well on a variety of natural language tasks without the need for extensive feature engineering. Baichuan 2 has shown significant improvements in various fields such as medicine and Chinese law compared to other models like ChatGLM and LLaMA.'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(query)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llangchainhf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
